#!/bin/bash

rm -rf ~/.triton/cache/*
rm -rf dump_dir/*
input_file=$1

export TRITON_DEBUG=1
export TRITON_DEVICE_PRINT=1
export TRITON_ALWAYS_COMPILE=1
export TRITON_LLVM_DEBUG_ONLY="triton-offset-analysis,triton-unstructure-converter"
export TRITON_REPRODUCER_PATH=./dump_dir
export TRITON_KERNEL_DUMP=1
export TRITON_DUMP_DIR=./dump_dir

# LOG_FILE="run_all_test.log"

# rm -f $LOG_FILE
run_and_log() {
    local cmd="$1"
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] Executing: $cmd" | tee $LOG_FILE -a 
    $cmd 2>&1 | tee -a $LOG_FILE
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] Command finish: $cmd" | tee $LOG_FILE -a
    echo -e "\n" >> $LOG_FILE 2>&1
}


# run_and_log "python3  causal_conv1d_fwd_kernel/*.py"
# run_and_log "python3  causal_conv1d_states_fwd_kernel/*.py"
# run_and_log "python3  causal_conv1d_update_kernel/*.py"
# run_and_log "python3  chunk_bwd_kernel_dqkwg/*.py"
# run_and_log "python3  chunk_bwd_kernel_dv_local/*.py"
# run_and_log "python3  chunk_fwd_kernel_o/*.py"
# run_and_log "python3  chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64/*.py"
# run_and_log "python3  chunk_gated_delta_rule_fwd_kernel_h_blockdim64/*.py"
# run_and_log "python3  chunk_local_cumsum_vector_kernel/*.py"
# run_and_log "python3  chunk_scaled_dot_kkt_fwd_kernel/*.py"
# run_and_log "python3  chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter/*.py"
# run_and_log "python3  fused_recurrent_gated_delta_rule_fwd_kernel/*.py"
# run_and_log "python3  merge_16x16_to_32x32_inverse_kernel/*.py"
# run_and_log "python3  merge_16x16_to_64x64_inverse_kernel/*.py"
# run_and_log "python3  prepare_wy_repr_bwd_kernel/*.py"
# run_and_log "python3  recompute_w_u_fwd_kernel/*.py"

# python3 precision_test.py >> $LOG_FILE 2>&1

### precision error
# run_and_log "python3  solve_tril_16x16_kernel/*.py"
### all pass
# run_and_log "python3  chunk_local_cumsum_scalar_kernel/*.py"
# run_and_log "python3  chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra/*.py"
# run_and_log "python3  l2norm_fwd_kernel/*.py"
# run_and_log "python3  layer_norm_gated_fwd_kernel/*.py"
# echo "All commands executed! Log saved to: $LOG_FILE"

function run_precision_for_all() {
    echo "Using BiSheng Install Path: $BISHENG_INSTALL_PATH"
    echo "Triton Ascend: "
    pip show triton-ascend
    file_array=("causal_conv1d_fwd_kernel" \
                "causal_conv1d_states_fwd_kernel" \
                "causal_conv1d_update_kernel" \
		"chunk_bwd_kernel_dqkwg" \
                "chunk_bwd_kernel_dv_local" \
                "chunk_fwd_kernel_o" \
                "chunk_gated_delta_rule_bwd_kernel_dhu_blockdim64" \
                "chunk_gated_delta_rule_fwd_kernel_h_blockdim64" \
                "chunk_local_cumsum_scalar_kernel" \
                "chunk_local_cumsum_vector_kernel" \
		"chunk_scaled_dot_kkt_fwd_kernel" \
                "chunk_scaled_dot_kkt_fwd_kernel_intra_sub_inter" \
                "chunk_scaled_dot_kkt_fwd_kernel_intra_sub_intra" \
                "fused_recurrent_gated_delta_rule_fwd_kernel" \
                "l2norm_fwd_kernel" \
                "layer_norm_gated_fwd_kernel" \
                "merge_16x16_to_32x32_inverse_kernel" \
                "merge_16x16_to_64x64_inverse_kernel" \
		"prepare_wy_repr_bwd_kernel" \
                "recompute_w_u_fwd_kernel" \
                "solve_tril_16x16_kernel" \
                "causal_conv1d_bwd_kernel" \
                "layer_norm_gated_bwd_kernel" \
                "layer_norm_gated_bwd_kernel1" )
    mkdir -p ./dump
    
    failed_files=()
    for filename in "${file_array[@]}"; do
        log_file="./dump/${filename}.log"
        echo "[$(date +'%Y-%m-%d %H:%M:%S')] Executing: $filename"
        python3 precision_test.py "$filename" 2>&1 | tee "$log_file"
        echo "[$(date +'%Y-%m-%d %H:%M:%S')] Finished: $filename"
        if [ ${PIPESTATUS[0]} -ne 0 ] || grep -q -i "AssertionError" "$log_file"; then
            failed_files+=("$filename")
        fi
    done
    
    echo "failed: ${failed_files[@]}"
}
run_precision_for_all
